library(gam)
library(foreign)

load("data/SMRPreplication.RData")

## Function to estimate prob of treatment: 
pscoreEstimate <- function(dat, id.var = "id", covs, exact.v, n.iter = 1, verbose = 100){
  store.p <- matrix(NA, n.iter, nrow(dat))
  store.p[, 1] <- 0.5

  for(iii in 1:n.iter){
    if(is.numeric(verbose) & ((iii %% verbose) == 0)){cat("Iteration ", iii, " of ", n.iter, "\n", sep ="")}
    sb1 <- seqblock1(id.vars = id.var, id.vals = dat[1, id.var], exact.vars = exact.v, exact.vals = dat[1, exact.v], covar.vars = covs, covar.vals = dat[1, covs], tr.names = c("anagrams", "memory"), assg.prob.kfac = 5, file.name = "~/Desktop/sbout.RData", verbose = FALSE)
  
    for(n.idx in 2:nrow(dat)){
      sb2k <- seqblock2k(object = "~/Desktop/sbout.RData", id.vals = dat[n.idx, id.var], exact.vals = dat[n.idx, exact.v], covar.vals = dat[n.idx, covs], file.name = "~/Desktop/sbout.RData", verbose = FALSE)
      ## store prob of getting "memory" condition:
      store.p[iii, n.idx] <- sb2k$p[which(sb2k$tr.sort == "memory")]
    }
  }
  return(round(store.p, 4))
}

## Define variables;
ccc <- c("pssi1_tot", "age", "phq1tot", "fas1tot", "stroopINTtscore1", "lm1_t1recallSCALED", "lm2_t1recallSCALED")
eee <- c("gender", "phq1mdddx")

## Run function above (optional):
#ps.mat <- pscoreEstimate(x[x$withdrawn == "no", ], covs = ccc, exact.v = eee, n.iter = 10000, verbose = 200)
#save(ps.mat, file = "data/pscoresFull10k.RData")
## Load output from one run of pscoreEstimate(), 10^4 iterations:
load("data/pscoresFull10k.RData")

# pedagogical xtable of pscore estimation:
store.p <- ps.mat[1:10, 1:8]
row.names(store.p) <- paste("replicate", 1:nrow(store.p), sep = "")
colnames(store.p) <- paste("unit", 1:ncol(store.p), sep = "")
library(xtable)
out.tex <- xtable(store.p, align=c("l", rep("c",ncol(store.p))))
#print(out.tex, file = "pscoresFULL.tex")
rm(store.p, out.tex)

## estimate each obs p-score by mean probability of treatment:
ps.vec <- apply(ps.mat, 2, mean)
summary(ps.vec)
## all \approx 0.5  

## using these probabilities of treatment:
source("code/calcAIPWest.R")
source("code/teEstimate.R")
diff.mean <- function(outcomes, dat=smrp.no){
  m.ci <- matrix(NA, length(outcomes), 3)
  for(i in 1:length(outcomes)){
    test.out <- t.test(dat[dat[,"condition"] == "memory", outcomes[i]], dat[dat[, "condition"] == "anagrams", outcomes[i]])
    m.ci[i, 1] <- test.out$estimate[1]-test.out$estimate[2]
    m.ci[i, 2:3] <- c(test.out$conf.int[1], test.out$conf.int[2])
  }
  return(m.ci)
}

outccc <- c("pssi2_tot", "pcl2tot", "phq2mdddx", "phq2tot", "riq2rum")
smrp.no <- x[x$withdrawn == "no", ]

## Estimate treatment effects (optional -- load final object below):
## set.seed(202)
## out.ests <- teEstimate(smrp.no, covs=ccc, exact.v=eee, n.iter=500, outcomes = outccc)
## out.ests.cr <- teEstimate(smrp.no, covs=ccc, exact.v=eee, n.iter=500, outcomes = outccc, assg = "cr")
## save(out.ests, file = "data/estsSBfull.RData")
## save(out.ests.cr, file = "data/estsCRfull.RData")
load("data/estsSBfull.RData")
load("data/estsCRfull.RData")

###### Compare Diff Means/AIPW on blocked data to complete randomization
plotdens2 <- function(outc, outname, xl, yl, zero = TRUE){
 plot(density(out.ests$ests.dm.mat[, outc]), xlim = xl, ylim = yl, xlab = paste(outname, ", null Tr Effect Estimates", sep=""), main = "")
 par(new = TRUE)
 plot(density(out.ests$ests.aipw.mat[, outc]), xlim = xl, ylim = yl, xlab = "", ylab="", lwd = 2, main = "")
 if(zero == TRUE){abline(v=0, col="grey", lty=2)}
}

## For all 5 outcomes:
# par(mfrow=c(3,2))
## For 2 illustrative outcomes:
par(mfrow = c(1,2))
plotdens2("pssi2_tot", "PTSD symptoms", xl = c(-10,10), yl = c(0, .35))
plotdens2("pcl2tot", "PTSD (self)", xl = c(-12,12), yl = c(0, .16))
plotdens2("phq2mdddx", "Depression (binary)", xl = c(-.8,.8), yl = c(0, 5))
plotdens2("phq2tot", "Depression (self)", xl = c(-6,6), yl = c(0, .5))
plotdens2("riq2rum", "Rumination", xl = c(-.8, .8), yl = c(0, 4))

## "less than one-fourth of an SD":
library(psych)
describe(out.ests[[1]])[,3]/describe(out.ests[[1]])[,4]
describe(out.ests[[2]])[,3]/describe(out.ests[[2]])[,4]

## Plot RMSE of SB ipw versus CR ipw:
outcomes <- c("pssi2_tot", "pcl2tot", "phq2mdddx", "phq2tot", "riq2rum")
rmse.sb <- data.frame(out.ests$rmse)
rmse.sb$iter <- 1:nrow(rmse.sb)
tmp <- reshape(rmse.sb, v.names = "rmse", varying = outcomes, direction = "long", idvar = "iter")
tmp$time <- outcomes[tmp$time]
names(tmp)[names(tmp)=="time"]  <- "outcome"
r.sb.long <- tmp 
##
rmse.cr <- data.frame(out.ests.cr$rmse)
rmse.cr$iter <- 1:nrow(rmse.cr)
tmp <- reshape(rmse.cr, v.names = "rmse", varying = outcomes, direction = "long", idvar = "iter")
tmp$time <- outcomes[tmp$time]
names(tmp)[names(tmp)=="time"]  <- "outcome"
r.cr.long <- tmp
rm(tmp)

## optional:  multiply riq2rum and phq2mdddx by k to make scales more comparable:
k <- 5
r.sb.long$rmse[r.sb.long$outcome == "riq2rum"] <- k*r.sb.long$rmse[r.sb.long$outcome == "riq2rum"]
r.sb.long$rmse[r.sb.long$outcome == "phq2mdddx"] <- k*r.sb.long$rmse[r.sb.long$outcome == "phq2mdddx"]
r.cr.long$rmse[r.cr.long$outcome == "riq2rum"] <- k*r.cr.long$rmse[r.cr.long$outcome == "riq2rum"]
r.cr.long$rmse[r.cr.long$outcome == "phq2mdddx"] <- k*r.cr.long$rmse[r.cr.long$outcome == "phq2mdddx"]
rm(k)

outcome.names <- c("PTSD 
                   symptoms", "PTSD 
                   (self)", "Depression 
                   (binary)", "Depression 
                   (self)", "Rumination")
boxwxx <- .08
var.space <- .3
barspace <- .045
xl <- c(0,14)

## Figure 9 (left):
par(mfrow=c(1,1))
boxplot(rmse ~ outcome, data=r.sb.long, horizontal = TRUE, xlim = c(0,1), ylim = xl, at = ((1:length(outcomes))+((-floor(length(outcomes)/2)):floor(length(outcomes)/2))*var.space)/(length(outcomes)+1)+barspace, boxwex = boxwxx, axes = FALSE, xlab = "RMSE")
boxplot(rmse ~ outcome, data=r.cr.long, horizontal = TRUE, xlim = c(0,1), ylim = xl, at = ((1:length(outcomes))+((-floor(length(outcomes)/2)):floor(length(outcomes)/2))*var.space)/(length(outcomes)+1)-barspace, boxwex = boxwxx, add = TRUE, axes = FALSE)
axis(1)
## (label order figured out by hand)
axis(2, at = ((1:length(outcomes))+((-floor(length(outcomes)/2)):floor(length(outcomes)/2))*var.space)/(length(outcomes)+1), labels = outcome.names[c(2, 3, 4, 1, 5)], las = 1, cex.axis=.7)

rm(outcomes, outcome.names, boxwxx, var.space, barspace, xl)
rm(r.cr.long, r.sb.long, rmse.cr, rmse.sb, ccc, eee, outccc)
